#!/usr/bin/env python3

from lxml import etree
from bs4 import BeautifulSoup
import struct
import sys
import zipfile

class FormatErrorException(Exception):
    """ AXML File Format Error  """
    def __init__(self, message=None):
        self.message = message

class TYPE:
    NULL = 0
    REFERENCE = 1
    ATTRIBUTE = 2
    STRING = 3
    FLOAT = 4
    DIMENSION = 5
    FRACTION = 6
    FIST_INT = 16
    INT_DEC = 16
    INT_HEX = 17
    INT_BOOLEAN = 18
    FIRST_COLOR_INT = 28
    INT_COLOR_ARGB8 = 28
    INT_COLOR_RGB8 = 29
    INT_COLOR_ARGB4 = 30
    INT_COLOR_RGB4 = 31
    LAST_COLOR_INT = 31
    LAST_INT = 31

class COMPLEX:
    UNIT_PX = 0
    UNIT_DIP = 1
    UNIT_SP = 2
    UNIT_PT = 3
    UNIT_IN = 4
    UNIT_MM = 5
    UNIT_SHIFT = 0,
    UNIT_MASK = 15
    UNIT_FRACTION = 0
    UNIT_FRACTION_PARENT = 1
    UNIT_RADIX_23p0 = 0
    UNIT_RADIX_16p7 = 1
    UNIT_RADIX_8p15 = 2
    UNIT_RADIX_0p23 = 3
    UNIT_RADIX_SHIFT = 4
    UNIT_MASK = 3
    UNIT_MANTISSA_SHIFT = 8
    UNIT_MANTISSA_MASK = 0xFFFFFF
    

class AXMLParser:
    """ parse axml file format  """
    def __init__(self, buffer):
        self.buffer = buffer
        self.file_size = None
        self.strings = None
        self.resources = None
        self._position = None
        self.string_pool = None
        self.resourceids = None
        self.elements = []
        self.root = None
        self.namespaces = {} # namespace uri to prefix
        self.tmp_namespace = {}     # prefix to namespace uri for tmp

    def check_magic(self):
        if self.buffer[:4] != b'\x03\x00\x08\x00':
            raise FormatErrorException('The magic number must be 0x03000800')

    def read_file_size(self):
        file_size_bytes = self.buffer[4:8]
        self.file_size = int.from_bytes(file_size_bytes, byteorder='little')

    def read_string_chunk(self):
        start = 8
        self._position = start
        chunk_type = self.buffer[self._position : self._position + 4]
        if chunk_type != b'\x01\x00\x1c\x00':
            raise FormatErrorException('The String Chunk ChunkType must be 0x01001c00, not ' + str(chunk_type))
        self._position += 4
        chunk_size = self._parse_u32int()
        string_count = self._parse_u32int()
        style_count = self._parse_u32int()
        self._position += 4 # pass Unknow
        string_pool_offset = self._parse_u32int()
        style_pool_offset = self._parse_u32int()
        string_offsets = [self._parse_u32int() for x in range(string_count)]
        style_offsets = [self._parse_u32int() for x in range(style_count)]
        str_pool_start = self._position
        self.string_pool = [self._parse_str(str_pool_start + str_offset) for str_offset in string_offsets]
        self._position = start + chunk_size

    def read_resourceid_chunk(self):
        chunk_type = self.buffer[self._position: self._position + 4]
        if chunk_type != b'\x80\x01\x08\x00':
            raise FormatErrorException('The ResourceId Chunk ChunkType must be 0x80010800')
        self._position += 4
        chunk_size = self._parse_u32int()
        self.resourceids = [hex(self._parse_u32int()) for x in range(int(chunk_size / 4) - 2)]

    def read_content_chunk(self):
        end = False
        while not end:
            chunk_type = self.buffer[self._position: self._position + 4]
            self._position += 4
            if chunk_type == b'\x00\x01\x10\x00':
                self._parse_start_name_space_chunk()
            elif chunk_type == b'\x02\x01\x10\x00':
                self._parse_start_tag_chunk()
            elif chunk_type == b'\x03\x01\x10\x00':
                self._parse_end_tag_chunk()
            elif chunk_type == b'\x04\x01\x10\x00':
                self._parse_text_chunk()
            elif chunk_type == b'\x01\x01\x10\x00':
                self._parse_end_name_space_chunk()
            elif len(chunk_type) == 0:
                end = True
            else:
                raise FormatErrorException(f"chunk_type: {chunk_type} not valid chunk_type")

    def _parse_text_chunk(self):
        chunk_size = self._parse_u32int()
        line_number = self._parse_u32int()
        self._position += 4
        name = self._parse_u32int()
        name_str = self.string_pool[name]
        self._position += 8
        self.elements[-1].text = name_str
        print(f'name: {name_str}')
        
    def _parse_end_tag_chunk(self):
        chunk_size = self._parse_u32int()
        line_number = self._parse_u32int()
        self._position += 4
        namespace_uri = self._parse_u32int()
        namespace_uri_str = self.string_pool[namespace_uri] if namespace_uri >= 0 else None
        name = self._parse_u32int()
        name_str = self.string_pool[name] if name >= 0 else None
        self.elements = self.elements[:-1]
        
    def _parse_start_name_space_chunk(self):
        chunk_size = self._parse_u32int()
        line_number = self._parse_u32int()
        self._position += 4
        prefix = self._parse_u32int()
        prefix_str = self.string_pool[prefix]
        uri = self._parse_u32int()
        uri_str = self.string_pool[uri]
        self.tmp_namespace[prefix_str] = uri_str
        self.namespaces[uri_str] = prefix_str
        

    def _parse_end_name_space_chunk(self):
        chunk_size = self._parse_u32int()
        line_number = self._parse_u32int()
        self._position += 4
        prefix = self._parse_u32int()
        uri = self._parse_u32int()
        prefix_str = self.string_pool[prefix]
        uri_str = self.string_pool[uri]

    def _parse_start_tag_chunk(self):
        chunk_size = self._parse_u32int()
        line_number = self._parse_u32int()
        self._position += 4
        namespace_uri = self._parse_u32int()
        namespace_uri_str = self.string_pool[namespace_uri] if namespace_uri >= 0 else None
        name = self._parse_u32int()
        name_str = self.string_pool[name]
        flags = self._parse_u32int()
        if flags != 1310740:
            print(f'flags not valid: {flags}')
        attribute_count = self._parse_u32int()
        class_attribute = self._parse_u32int()
        attributes = {}
        for x in range(attribute_count):
            uri = self._parse_u32int()
            uri_str = self.string_pool[uri] if uri >= 0 else None
            attribute_name = self._parse_u32int()
            attribute_name_str = self.string_pool[attribute_name] if attribute_name >= 0 else None
            value = self._parse_u32int()
            value_str = self.string_pool[value] if value >= 0 else None
            type = self.buffer[self._position: self._position + 4]
            data = self.buffer[self._position + 4: self._position + 8]
            self._position += 8
            name_with_prefix = "{" + uri_str + "}" + attribute_name_str if uri_str else attribute_name_str
            if value_str:
                attributes[name_with_prefix] = value_str
            else:
                attributes[name_with_prefix] = AXMLParser._parse_attrib_value(type, data) or "unknow"
        if len(self.elements) == 0:
            element = etree.Element(name_str, nsmap=self.tmp_namespace, attrib=attributes)
            self.root = element
        else:
            element = etree.SubElement(self.elements[-1], name_str, nsmap=self.tmp_namespace, attrib=attributes)
        self.tmp_namespace = {}
        self.elements.append(element)

    @staticmethod
    def _parse_attrib_value(_type, data):
        _type = int.from_bytes(_type, byteorder='little', signed=True)
        _type = _type >> 24
        data = int.from_bytes(data, byteorder='little', signed=True)
        dimension_units = ['px', 'dip', 'sp', 'pt', 'in', 'mm', '', '']
        fraction_units =  ['%',  '%p',  '',   '',   '',   '',   '', '']
        type_case = {
            TYPE.NULL: '',
            TYPE.ATTRIBUTE: lambda x : "?%s%08X" % (AXMLParser._get_package(x), x),
            TYPE.REFERENCE: lambda x : "@%s%08X" % (AXMLParser._get_package(x), x),
            TYPE.FLOAT: lambda x: str(AXMLParser._int32_to_float(x)),
            TYPE.INT_HEX: lambda x: "0x%08X" % x,
            TYPE.INT_BOOLEAN: lambda x: "true" if x != 0 else "false",
            TYPE.DIMENSION: lambda x: str(AXMLParser._complex_to_float(x)) + dimension_units[data & COMPLEX.UNIT_MASK],
            TYPE.FRACTION: lambda x: str(AXMLParser._complex_to_float(x)) + fraction_units[data & COMPLEX.UNIT_MASK]
        }
        if _type in type_case:
            case = type_case[_type]
            if callable(case):
                return case(data)
            else:
                return case
        if _type >= TYPE.FIRST_COLOR_INT and _type <= TYPE.LAST_COLOR_INT:
            return "#%08X" % data
        if _type >= TYPE.FIST_INT and _type <= TYPE.LAST_INT:
            return str(data)
        return "<0x%X, type 0x%02X>" % (data, _type)

    @staticmethod
    def _complex_to_float(complex_value):
        RADIX_MULTS = [0.00390625,3.051758E-5,1.192093E-7,4.656613E-10]
        return (complex_value & 0xFFFFFF00) * RADIX_MULTS[(complex_value >> 4) & 3]
        
    @staticmethod
    def _int32_to_float(data):
        return struct.unpack('<f', struct.pack('<i', data))[0]
    @staticmethod
    def _attrib_value_attribute(data):
        return "?%s%08X" % (AXMLParser._get_package(data), data)

    @staticmethod
    def _get_package(id):
        if id >> 24 == 1:
            return 'android/'
        return 'ref/'
            
    def _parse_str(self, start):
        end = start + 2
        size = int.from_bytes(self.buffer[start: end], byteorder='little') * 2
        str_values = b'\xff\xfe' + self.buffer[end: end + size]     # little-endian
        return str_values.decode('utf-16')
            
    def _parse_u32int(self):
        values_bytes = self.buffer[self._position: self._position + 4]
        self._position += 4
        return int.from_bytes(values_bytes, byteorder='little', signed=True)
        
        
if __name__ == '__main__':
    file_path = sys.argv[1]
    buffer = None
    if file_path.endswith('.apk'):
        with zipfile.ZipFile(file_path) as myzip:
            with myzip.open('AndroidManifest.xml', 'r') as myfile:
                buffer = myfile.read()
    else:
        with open(file_path, 'rb') as file:
            buffer = file.read()
    if buffer:
        parser = AXMLParser(buffer)
        parser.check_magic()
        parser.read_file_size()
        parser.read_string_chunk()
        parser.read_resourceid_chunk()
        parser.read_content_chunk()
        print(BeautifulSoup(etree.tostring(parser.root), 'xml').prettify())
